package ssau.knyazev.fft.modules.cuda;

import jcuda.Pointer;

import jcuda.Sizeof;
import jcuda.driver.CUcontext;
import jcuda.driver.CUctx_flags;
import jcuda.driver.CUdevice;
import jcuda.driver.CUdeviceptr;
import jcuda.driver.CUfunction;
import jcuda.driver.CUmodule;
import jcuda.driver.JCudaDriver;
import ssau.knyazev.common.CudaCRuntimeException;
import ssau.knyazev.fft.common.FFTConst;
import ssau.knyazev.modules.CudaCompiler;
import ssau.knyazev.modules.Generator;
import ssau.knyazev.modules.Printer;

public class CudaCFFTModule {

	private static CudaCFFTModule instance = null;
	private CUdevice device = null;
	private CUcontext context = null;
	private CUfunction fftFunction = null;
	private CUfunction fftMirrorFunction = null;
	private CUfunction wlistFunction = null;
	private CUfunction invRowsFunction = null;
	private CUfunction invColsFunction = null;
	private CUfunction inverseFunction = null;
	private CUfunction convertToCComplexFunction = null;
	private CUfunction convertToComplexFunction = null;
	private static long sum = 0;
	private static int count = 10;
	private int blockSize = 0;
	private int gridSize = 0;
	
	private CudaCFFTModule() throws CudaCRuntimeException{
		JCudaDriver.setExceptionsEnabled(true);
		JCudaDriver.cuInit(0);
		device = new CUdevice();
		JCudaDriver.cuDeviceGet(device, 0);
		context = new CUcontext();
		JCudaDriver.cuCtxCreate(context, CUctx_flags.CU_CTX_BLOCKING_SYNC, device);
		String moduleName = CudaCompiler.preparePtxFile(FFTConst.CUDA_CCOMPLEX_MODULE_SRC, false);
		CUmodule module = new CUmodule();
		JCudaDriver.cuModuleLoad(module, moduleName);
		fftFunction = new CUfunction();
		JCudaDriver.cuModuleGetFunction(fftFunction, module, FFTConst.CUDA_FFT_FUNCTION_NAME);
		wlistFunction = new CUfunction();
		JCudaDriver.cuModuleGetFunction(wlistFunction, module, FFTConst.CUDA_WLIST_FUNCTION_NAME);
		invColsFunction = new CUfunction();
		JCudaDriver.cuModuleGetFunction(invColsFunction, module, "inverseCols");
		invRowsFunction = new CUfunction();
		JCudaDriver.cuModuleGetFunction(invRowsFunction, module, "inverseRows");
		inverseFunction = new CUfunction();
		JCudaDriver.cuModuleGetFunction(inverseFunction, module, "inverse");
		fftMirrorFunction = new CUfunction();
		JCudaDriver.cuModuleGetFunction(fftMirrorFunction, module, "fftMirror");
		convertToCComplexFunction = new CUfunction();
		JCudaDriver.cuModuleGetFunction(convertToCComplexFunction, module, "convertToCComplex");
		convertToComplexFunction = new CUfunction();
		JCudaDriver.cuModuleGetFunction(convertToComplexFunction, module, "convertToComplex");
	}
	
	public static CudaCFFTModule createInstance() throws CudaCRuntimeException{
		if (instance != null){
			return instance;
		} else {
			instance = new CudaCFFTModule();
			return instance;
		}
	}
	
	public float[][] fft(float[][] in) throws CudaCRuntimeException{
		int length = in.length;
		if (in != null && in.length == in[0].length){
			in = transformToComplexMatrix(in);
		}
		float[] arr = toVector(in);
		blockSize = 16;
		gridSize = (length > blockSize) ? length/blockSize : 1;
		
		CUdeviceptr wilistPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(wilistPtr, length * 4* Sizeof.FLOAT);
		CUdeviceptr wjlistPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(wjlistPtr, length * 4* Sizeof.FLOAT);
		
		CUdeviceptr incomingPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(incomingPtr, arr.length * Sizeof.FLOAT);
		JCudaDriver.cuMemcpyHtoD(incomingPtr, Pointer.to(arr), arr.length * Sizeof.FLOAT);
		
		CUdeviceptr srcPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(srcPtr, arr.length *4 * Sizeof.FLOAT);
		
		//START FFT
		long time = System.nanoTime();
		int size = 0;
		for (int x = length; x > 1; x = x >> 1){
			size++;
		}
		convertToCComplex(in.length, incomingPtr, srcPtr);
		inverse(in.length, size-1, srcPtr);
		calculateWList(length, wilistPtr, wjlistPtr);
		
		//System.out.println(length);
		for (int iteration = 0; iteration < size; iteration++){
			long del = System.nanoTime();
			Pointer kernelParameters = Pointer.to(
					Pointer.to(new int[] { length }), 
					Pointer.to(srcPtr),
					Pointer.to(wilistPtr),
					Pointer.to(wjlistPtr),
					Pointer.to(new int[] { iteration })
					);
			JCudaDriver.cuLaunchKernel(fftFunction, 
					gridSize, gridSize, 1,
					blockSize, blockSize, 1,
					0, null,
					kernelParameters, null
			);
			JCudaDriver.cuCtxSynchronize();
			del = System.nanoTime() - del;
			System.out.println(iteration + ": " + del);
		}
		convertToComplex(length, incomingPtr, srcPtr);
		time = System.nanoTime() - time;
		sum += time;
		//FINISH FFT
		
		JCudaDriver.cuMemcpyDtoH(Pointer.to(arr), incomingPtr, arr.length * Sizeof.FLOAT);
		
		in = toMatrix(arr, in.length);
		
		JCudaDriver.cuMemFree(wilistPtr);
		JCudaDriver.cuMemFree(incomingPtr);
		JCudaDriver.cuMemFree(srcPtr);
		return in;
	}
	
	public float[][] fftMirrorMode(float[][] in) throws CudaCRuntimeException{
		int length = in.length;
		if (in != null && in.length == in[0].length){
			in = transformToComplexMatrix(in);
		}
		float[] arr = toVector(in);
		
		CUdeviceptr wilistPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(wilistPtr, length * 4* Sizeof.FLOAT);
		CUdeviceptr wjlistPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(wjlistPtr, length * 4* Sizeof.FLOAT);
		
		CUdeviceptr incomingPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(incomingPtr, arr.length * Sizeof.FLOAT);
		JCudaDriver.cuMemcpyHtoD(incomingPtr, Pointer.to(arr), arr.length * Sizeof.FLOAT);
		
		CUdeviceptr srcPtr = new CUdeviceptr();
		JCudaDriver.cuMemAlloc(srcPtr, arr.length *4 * Sizeof.FLOAT);
		
		//START FFT
		long time = System.nanoTime();
		int size = 0;
		for (int x = length; x > 1; x = x >> 1){
			size++;
		}
		convertToCComplex(in.length, incomingPtr, srcPtr);
		inverse(in.length, size-1, srcPtr);
		calculateWList(length, wilistPtr, wjlistPtr);
		
		for (int iteration = 0; iteration < size; iteration++){
			Pointer kernelParameters = Pointer.to(
					Pointer.to(new int[] { length }), 
					Pointer.to(srcPtr),
					Pointer.to(wilistPtr),
					Pointer.to(wjlistPtr),
					Pointer.to(new int[] { iteration })
					);
			JCudaDriver.cuLaunchKernel(fftMirrorFunction, 
					8, 8, 1,
					8, 8, 1,
					0, null,
					kernelParameters, null
			);
			JCudaDriver.cuCtxSynchronize();
		}
		convertToComplex(length, incomingPtr, srcPtr);
		time = System.nanoTime() - time;
		sum += time;
		//FINISH FFT
		
		JCudaDriver.cuMemcpyDtoH(Pointer.to(arr), incomingPtr, arr.length * Sizeof.FLOAT);
		
		in = toMatrix(arr, in.length);
		
		JCudaDriver.cuMemFree(wilistPtr);
		JCudaDriver.cuMemFree(incomingPtr);
		JCudaDriver.cuMemFree(srcPtr);
		return in;
	}
	
	private void inverse(int length, CUdeviceptr srcPtr){
		Pointer kernelParameters = Pointer.to(
				Pointer.to(new int[] { length}), 
				Pointer.to(srcPtr)
				);
		JCudaDriver.cuLaunchKernel(invRowsFunction, 
				8, 8, 1,
				8, 8, 1,
				0, null,
				kernelParameters, null
		);
		JCudaDriver.cuCtxSynchronize();
		kernelParameters = Pointer.to(
				Pointer.to(new int[] { length}), 
				Pointer.to(srcPtr)
				);
		JCudaDriver.cuLaunchKernel(invColsFunction, 
				8, 8, 1,
				8, 8, 1,
				0, null,
				kernelParameters, null
		);
	}
	
	private void inverse(int length, int power, CUdeviceptr srcPtr){
		Pointer kernelParameters = Pointer.to(
				Pointer.to(new int[] { length}),
				Pointer.to(new int[] { power}),
				Pointer.to(srcPtr)
				);
		JCudaDriver.cuLaunchKernel(inverseFunction, 
				gridSize, gridSize, 1,
				blockSize, blockSize, 1,
				0, null,
				kernelParameters, null
		);
		JCudaDriver.cuCtxSynchronize();
	}
	
	private void calculateWList(int length, CUdeviceptr wilistPtr, CUdeviceptr wjlistPtr){
		Pointer kernelParameters = Pointer.to(
				Pointer.to(new int[] { length }), 
				Pointer.to(wilistPtr),
				Pointer.to(wjlistPtr)
				);
		JCudaDriver.cuLaunchKernel(wlistFunction, 
				gridSize, gridSize, 1,
				blockSize, blockSize, 1,
				0, null,
				kernelParameters, null
		);
		JCudaDriver.cuCtxSynchronize();
	}
	
	private void convertToCComplex(int length, CUdeviceptr complexPtr, CUdeviceptr ccomplexPtr){
		Pointer kernelParameters = Pointer.to(
				Pointer.to(new int[] { length }), 
				Pointer.to(complexPtr),
				Pointer.to(ccomplexPtr)
				);
		JCudaDriver.cuLaunchKernel(convertToCComplexFunction, 
				gridSize, gridSize, 1,
				blockSize, blockSize, 1,
				0, null,
				kernelParameters, null
		);
		JCudaDriver.cuCtxSynchronize();
	}
	
	private void convertToComplex(int length, CUdeviceptr complexPtr, CUdeviceptr ccomplexPtr){
		Pointer kernelParameters = Pointer.to(
				Pointer.to(new int[] { length }), 
				Pointer.to(complexPtr),
				Pointer.to(ccomplexPtr)
				);
		JCudaDriver.cuLaunchKernel(convertToComplexFunction, 
				gridSize, gridSize, 1,
				blockSize, blockSize, 1,
				0, null,
				kernelParameters, null
		);
		JCudaDriver.cuCtxSynchronize();
	}
	
	protected static float[] toVector(float[][] src) {
		float[] res = new float[src.length * src[0].length];
		int k = 0;
		for (int i = 0; i < src.length; i++) {
			for (int j = 0; j < src[0].length; j++) {
				res[k] = src[i][j];
				k++;
			}
		}
		return res;
	}
	
	protected static float[][] toMatrix(float[] src, int rows) {
		float[][] res = new float[rows][src.length / rows];
		int k = 0;
		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < src.length / rows; j++) {
				res[i][j] = src[k];
				k++;
			}
		}
		return res;
	}
	
	protected static float[][] transformToComplexMatrix(float[][] src){
		float[][] res = new float[src.length][src[0].length*2];
		for (int i = 0; i < src.length; i++){
			for (int j = 0; j < src.length; j++){
				res[i][j*2] = src[i][j];
			}
		}
		return res;
	}
	
	public static void main(String[] args) {
		float[][] src = Generator.generateSimpleMatrix(4096);
		long time = 0;
		try {
			CudaCFFTModule core = createInstance();
			for (int i = 0; i < count; i++){
				src = core.fft(src);
			}
		} catch (CudaCRuntimeException e) {
			e.printStackTrace();
		}
		
		
		time = sum/count;
		System.out.println(time);
	}
	
}
